# -*- coding: utf-8 -*-
# created on 2016/11/30

from sympy import Abs
from mathsolver.functions.base import *
from mathsolver.functions.base.base import new_latex
from mathsolver.functions.fushu.base import fs_default_symbol
from sympy.abc import z, Z
from mathsolver.functions.fushu.base import Complex
from mathsolver.functions.fushu.basic import fs_search_multiple
from mathsolver.functions.mathematica.mathematicaSolve import MathematicaSolve
from mathsolver.functions.fushu.FuShu_BiaoZhunHua import FSNormalized


class FSSolveEq001(BaseFunction):
    """
    z = 5 * I - 3 * z
    :解复数方程：必须含z，方程可解
    """
    def solver(self, *args):
        eq = args[0].sympify()
        expr = eq[0] - eq[1]
        expr = expr.expand().simplify()
        symbols = expr.free_symbols
        assert (z in symbols) or (Z in symbols) or (sympify('z_1') in symbols) or (sympify('z_2') in symbols) or (sympify('omega') in symbols)
        symbol = fs_default_symbol(expr)
        try:
            solutions = solve(expr, symbol)
            new_solutions = []
            for solution in solutions:
                solution = FSNormalized().solver(BaseComplexPoly(solution)).output[0].sympify()
                new_solutions.append(solution)
        except Exception:
            raise Exception("try fail")
        assert solutions
        answer = {}
        values = []
        values.extend(new_solutions)
        answer[symbol] = values
        self.steps.append(["", "解复数方程, 得"])
        self.steps.append(["", "或".join(["∴ %s = %s" % (new_latex(symbol), new_latex(v)) for v in values])])
        self.output.append(BaseComplexValue(answer))
        self.label.add("解复数方程")
        return self


class FSSolveEq002(BaseFunction):
    """
    若i(x + yi)=3+4i,x,y∈R,则复数x+yi的模是 ()
    :不包含z,Z
    """
    def solver(self, *args):
        eq = args[0].sympify()
        expr = eq[0] - eq[1]
        symbols = expr.free_symbols
        assert (z not in symbols) and (Z not in symbols)
        expr = expr.expand().simplify()
        expr = FSNormalized().solver(BaseComplexPoly(expr)).output[0].sympify()
        abs_childrens = get_all_child(expr, lambda xx: isinstance(xx, Abs))
        if abs_childrens:
            for Abs_child in abs_childrens:
                abs_value = Abs_child.args[0]
                abs_value_com = Complex(abs_value)
                abs_value_com_real = abs_value_com.real
                abs_value_com_imag = abs_value_com.imag
                expr = expr.subs({Abs_child: sqrt(abs_value_com_real ** 2 + abs_value_com_imag ** 2)})
        com = Complex(expr)
        com_real = com.real
        com_imag = com.imag
        self.output.append(BaseEqs([[com_real, S.Zero], [com_imag, S.Zero]]))
        return self


class FSSolveEq003(BaseFunction):
    """
    已知\\frac{\\overline{z}}{1 + i} = 2 + i,则复数z =( )
    """
    def solver(self, *args):
        left, right = args[0].sympify()
        expr = (left - right).expand().simplify()
        symbol = fs_default_symbol(expr)
        sym_value = sympify('a') + sympify('b') * sympify('I')
        sym_com = Complex(sym_value)
        sym_real = sym_com.real
        sym_imag = sym_com.imag
        self.steps.append(["", "解方程: %s" % (BaseEq([left, right]).printing())])
        conjugate_childrens = get_all_child(expr, lambda xx: xx.__class__.__name__ == 'conjugate')
        new_expr = expr
        if conjugate_childrens:
            for conjugate_child in conjugate_childrens:
                new_expr = new_expr.subs({conjugate_child: sym_real - sym_imag * sympify('I')})
        abs_childrens = get_all_child(expr, lambda xx: isinstance(xx, Abs))
        if abs_childrens:
            for Abs_child in abs_childrens:
                new_expr = new_expr.subs({Abs_child: sqrt(sym_real ** 2 + sym_imag ** 2)})
        new_expr = new_expr.subs({symbol: sym_value})
        new_expr = new_expr.expand().simplify()
        new_expr_com = Complex(new_expr)
        new_expr_real = new_expr_com.real
        new_expr_imag = new_expr_com.imag
        eqs = [[new_expr_real, S.Zero], [new_expr_imag, S.Zero]]
        solutions = MathematicaSolve().solver(BaseEqs(eqs)).output[0].sympify()
        answers = fs_search_multiple(solutions, sym_value)
        sym_values = []
        for ans in answers:
            new_sym_value = sym_value
            new_sym_value = new_sym_value.subs(ans)
            new_sym_value = new_sym_value.expand().simplify()
            sym_values.append(new_sym_value)
            self.steps.append(["", "∴ 解得: %s" % " 或 ".join(["%s = %s" % (new_latex(symbol), new_latex(v)) for v in sym_values])])
        self.output.append(BaseComplexValue({symbol: sym_values}))
        self.label.add("解复数方程-待定系数法")
        return self


class FSSolveEq(BaseFunction):
    CLS = [FSSolveEq001, FSSolveEq002, FSSolveEq003]

    def solver(self, *args):
        assert len(args) == 1
        assert isinstance(args[0], (BaseComplexEq, BaseEq))
        known = self.known
        r = None
        for cl in FSSolveEq.CLS:
            try:
                new_known = dict(known)
                r = cl(known=new_known, verbose=True).solver(*args)
                break
            except Exception:
                pass
        if not r:
            raise 'try fail'
        return r


class FSSolveEqs(BaseFunction):
    """
    :FSSolveEqs().solver(BaseEqs([['conjugate(z)+z', '2'], ['(z-conjugate(z))*I', '2']]))
    :若复数{z_1} = 1 + {{i}},{z_2} = 3 -{{ i}},则{z_1} • {z_2} =  ()
    """

    def solver(self, *args):
        assert len(args) == 1
        eqs = args[0].sympify()
        assert len(eqs) > 1
        solutions = {}
        symbols = set([])
        for eq in eqs:
            expr = (eq[0] - eq[1]).expand().simplify()
            expr_symbols = expr.free_symbols
            symbols = symbols | expr_symbols
        symbols = list(symbols)
        symbol_values = {}
        for index, symbol in enumerate(symbols):
            real = str('a') + str(index + 1)
            imag = str('b') + str(index + 1)
            symbol_values[symbol] = sympify(real) + sympify(imag) * sympify('I')
            conjugate = str('conjugate') + str('(') + str(symbol) + str(')')
            conjugate = sympify(conjugate)
            symbol_values[conjugate] = sympify(real) - sympify(imag) * sympify('I')
            self.steps.append(["", "设%s = %s" % (new_latex(symbol), new_latex(symbol_values[symbol]))])
        feqs = []
        for eq in eqs:
            expr = (eq[0] - eq[1]).expand().simplify()
            expr = expr.subs(symbol_values)
            expr = expr.expand().simplify()
            expr = Complex(expr)
            expr_real = expr.real
            if expr_real != 0:
                feqs.append([expr_real, S.Zero])
            expr_imag = expr.imag - 0
            if expr_imag != 0:
                feqs.append([expr_imag, S.Zero])
        if len(feqs) == 1:
            self.steps.append(["", "%s" % (BaseEq(feqs[0]).printing())])
        else:
            self.steps.append(["", "%s" % (BaseEqs(feqs).printing())])
        feqs_solutions = MathematicaSolve().solver(BaseIneqs(feqs)).output[0].sympify()
        assert feqs_solutions
        self.steps.append(["", "解得: %s" % (BaseSymbolValue(feqs_solutions).printing())])
        for symbol in symbols:
            symbol_value = symbol_values[symbol]
            symbol_value_answers = fs_search_multiple(feqs_solutions, symbol_value)
            assert symbol_value_answers
            target_values = []
            for symbol_value_answer in symbol_value_answers:
                new_value = symbol_value
                new_value = new_value.subs(symbol_value_answer)
                new_value = new_value.expand().simplify()
                target_values.append(new_value)
            solutions[symbol] = target_values
        self.steps.append(["", "%s" % (BaseSymbolValue(solutions).printing())])
        self.output.append(BaseComplexValue(solutions))
        self.label.add("解复数方程组")
        return self


if __name__ == '__main__':
    pass
